from gxl_ai_utils.utils import utils_file
import random


class ICL_sampler:
    def __init__(self, wav_dict_path=None, text_dict_path=None):
        if not (wav_dict_path is None or text_dict_path is None):
            if isinstance(wav_dict_path, dict):
                self.wav_dict = wav_dict_path
                self.text_dict = text_dict_path
            elif isinstance(wav_dict_path, str):
                self.wav_dict = utils_file.load_dict_from_scp(wav_dict_path)
                self.text_dict = utils_file.load_dict_from_scp(text_dict_path)
            self.key_list = list(self.wav_dict.keys())
        else:
            self.wav_dict = wav_dict_path
            self.text_dict = text_dict_path


    def set_source(self, wav_dict_path, text_dict_path):
        if isinstance(wav_dict_path, dict):
            self.wav_dict = wav_dict_path
            self.text_dict = text_dict_path
        else:
            self.wav_dict = utils_file.load_dict_from_scp(wav_dict_path)
            self.text_dict = utils_file.load_dict_from_scp(text_dict_path)
        self.key_list = list(self.wav_dict.keys())

    def get_num_pair(self, num_pair):
        if num_pair == len(self.key_list):
            utils_file.logging_limit_print('get_num_pair: len(self.key_list) == num_pair')
            wav_path_list = []
            text_list = []
            for key in self.key_list:
                wav_path = self.wav_dict[key]
                text = self.text_dict[key]
                wav_path_list.append(wav_path)
                text_list.append(text)
            return wav_path_list, text_list

        key_list = []
        wav_path_list = []
        text_list = []
        index = 0
        while True:
            key = random.choice(list(self.wav_dict.keys()))
            if key in key_list:
                continue
            key_list.append(key)
            wav_path = self.wav_dict[key]
            text = self.text_dict[key]
            if len(text) < 1 or len(text) > 8:
                continue
            wav_path_list.append(wav_path)
            text_list.append(text)
            index += 1
            if index == num_pair:
                break
        return wav_path_list, text_list


def get_10_random_wenetspeech_pair():
    wav_path_list, text_list = global_sampler.get_num_pair(10)
    return wav_path_list, text_list


global_sampler = ICL_sampler(None, None)
